# main_beihang.py
"""
北航航线规划API服务入口
专门处理北航格式的航线规划请求
"""
import uvicorn
from fastapi import FastAPI, HTTPException, Request
import json
from datetime import datetime
import argparse
import matplotlib.pyplot as plt
from typing import Dict, Any
import logging
import os
import sys
import socket
from pathlib import Path

# 尝试导入psutil，如果不存在则提供替代方案
try:
    import psutil
    PSUTIL_AVAILABLE = True
except ImportError:
    PSUTIL_AVAILABLE = False

# 导入北航格式模型和转换器
from FastAPI.beihang_models import BeihangRoutePlanRequest, BeihangRoutePlanResponse
from FastAPI.beihang_converter import BeihangConverter
from FastAPI.planner_service import run_planning_service
from FastAPI.logger_config import get_logger

# --- 获取日志器 ---
logger = get_logger('beihang_api')

# --- FastAPI 应用配置 ---
app = FastAPI(
    title="北航航线规划API",
    description="专门处理北航格式的无人机任务规划API",
    version="1.3.0",
)

# 全局可视化设置
global_visualize_enabled = False

@app.get("/health", summary="健康检查")
async def health_check():
    """
    服务器健康检查接口
    """
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "service": "beihang_route_planning_api",
        "version": "1.3.0"
    }

def find_and_kill_process_on_port(port: int) -> bool:
    """
    查找并杀死占用指定端口的进程
    
    Args:
        port: 端口号
        
    Returns:
        bool: 是否成功杀死进程
    """
    # 首先尝试使用psutil
    if PSUTIL_AVAILABLE:
        try:
            # 使用更简单的方法查找占用端口的进程
            for proc in psutil.process_iter(['pid', 'name']):
                try:
                    # 获取进程的网络连接信息
                    connections = proc.connections()
                    if connections:
                        for conn in connections:
                            if conn.laddr and conn.laddr.port == port:
                                logger.info(f"发现占用端口 {port} 的进程: {proc.info['name']} (PID: {proc.info['pid']})")
                                print(f"发现占用端口 {port} 的进程: {proc.info['name']} (PID: {proc.info['pid']})")
                                proc.kill()
                                logger.info(f"已杀死进程 {proc.info['pid']}")
                                print(f"已杀死进程 {proc.info['pid']}")
                                return True
                except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                    continue
                except Exception as e:
                    logger.warning(f"检查进程 {proc.info.get('pid', 'unknown')} 时出错: {e}")
                    continue
        except Exception as e:
            logger.error(f"使用psutil查找进程时出错: {e}")
    
    # 如果psutil失败，尝试使用系统命令
    return kill_process_by_system_command(port)

def kill_process_by_system_command(port: int) -> bool:
    """
    使用系统命令杀死占用端口的进程
    
    Args:
        port: 端口号
        
    Returns:
        bool: 是否成功杀死进程
    """
    try:
        import platform
        system = platform.system()
        
        if system == 'Windows':
            # Windows系统使用netstat和taskkill
            import subprocess
            result = subprocess.run(['netstat', '-ano'], capture_output=True, text=True)
            if result.returncode == 0:
                lines = result.stdout.split('\n')
                for line in lines:
                    if f':{port}' in line and 'LISTENING' in line:
                        parts = line.strip().split()
                        if len(parts) >= 5:
                            pid = parts[-1]
                            print(f"发现占用端口 {port} 的进程PID: {pid}")
                            
                            # 杀死进程
                            kill_result = subprocess.run(['taskkill', '/F', '/PID', pid], capture_output=True, text=True)
                            if kill_result.returncode == 0:
                                logger.info(f"已杀死PID {pid} 的进程")
                                print(f"已杀死PID {pid} 的进程")
                                return True
                            else:
                                logger.warning(f"杀死PID {pid} 失败: {kill_result.stderr}")
                                print(f"杀死PID {pid} 失败")
        
        elif system == 'Linux':
            # Linux系统使用lsof和kill
            import subprocess
            result = subprocess.run(['lsof', '-ti', f':{port}'], capture_output=True, text=True)
            if result.returncode == 0:
                pid = result.stdout.strip()
                if pid:
                    print(f"发现占用端口 {port} 的进程PID: {pid}")
                    
                    # 杀死进程
                    kill_result = subprocess.run(['kill', '-9', pid], capture_output=True, text=True)
                    if kill_result.returncode == 0:
                        logger.info(f"已杀死PID {pid} 的进程")
                        print(f"已杀死PID {pid} 的进程")
                        return True
                    else:
                        logger.warning(f"杀死PID {pid} 失败")
                        print(f"杀死PID {pid} 失败")
        
        return False
        
    except Exception as e:
        logger.error(f"使用系统命令杀死进程时出错: {e}")
        print(f"使用系统命令杀死进程时出错: {e}")
        return False

def is_port_in_use(port: int) -> bool:
    """
    检查端口是否被占用
    
    Args:
        port: 端口号
        
    Returns:
        bool: 端口是否被占用
    """
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(1)
            result = s.connect_ex(('localhost', port))
            return result == 0
    except Exception:
        return False

@app.post("/plan_route_beihang", response_model=BeihangRoutePlanResponse, summary="北航格式任务规划")
async def plan_route_beihang(request: BeihangRoutePlanRequest, http_request: Request = None):
    global global_visualize_enabled
    """
    接收北航格式的任务规划请求，返回北航格式的规划结果。
    
    Args:
        request: 北航格式的任务规划请求
    """
    try:
        # 记录请求信息
        client_ip = http_request.client.host if http_request else "unknown"
        logger.info(f"[北航接口] 收到来自 {client_ip} 的规划请求")
        logger.info(f"[北航接口] 请求ID: {request.outBizCode}")
        logger.info(f"[北航接口] 路由类型: {request.routeType}")
        logger.info(f"[北航接口] 设备数量: {len(request.deviceInfos)}")
        logger.info(f"[北航接口] 可视化参数: {global_visualize_enabled}")
        print(f"[DEBUG] global_visualize_enabled = {global_visualize_enabled}")

        # 1. 转换北航格式为内部格式
        logger.info("[北航接口] 开始转换北航格式为内部格式...")
        converter = BeihangConverter()
        internal_request = converter.beihang_to_internal(request.model_dump())
        logger.info("[北航接口] 格式转换完成")
        
        # 2. 如果启用可视化，绘制环境数据
        if global_visualize_enabled:
            logger.info("[北航接口] 开始可视化环境数据...")
            visualize_environment_data(request.environmentalInfo.model_dump(), internal_request['config'])
        else:
            logger.info("[北航接口] 跳过环境数据可视化")

        # 3. 调用核心规划算法
        logger.info("[北航接口] 开始调用核心规划算法...")
        start_time = datetime.now()
        internal_response = run_planning_service(internal_request)
        end_time = datetime.now()
        algorithm_duration = (end_time - start_time).total_seconds()
        logger.info(f"[北航接口] 核心规划算法完成，耗时: {algorithm_duration:.2f}秒")
        
        # 4. 转换内部格式为北航格式
        logger.info("[北航接口] 开始转换内部格式为北航格式...")
        beihang_response = converter.internal_to_beihang({
            'status': 'success',
            'message': '任务处理成功',
            'data': {
                'trajectories': internal_response[0],
                'time_paths': internal_response[1]
            }
        }, request.model_dump())
        logger.info("[北航接口] 北航格式转换完成")
        
        # 5. 如果启用可视化，绘制轨迹结果
        if global_visualize_enabled:
            logger.info("[北航接口] 开始可视化规划结果...")
            visualize_trajectories(internal_response[0], internal_response[1], internal_request['config'])
        else:
            logger.info("[北航接口] 跳过轨迹结果可视化")
        
        # 6. 验证返回数据格式
        logger.info("[北航接口] 开始验证返回数据格式...")
        is_valid = validate_beihang_response_format(beihang_response)
        if not is_valid:
            logger.warning("[北航接口] 返回数据格式验证失败!")
        else:
            logger.info("[北航接口] 数据格式验证通过")
        
        # 记录成功响应信息
        route_count = len(beihang_response.get('routeInfos', []))
        logger.info(f"[北航接口] 请求处理成功，生成 {route_count} 条路径")
        logger.info(f"[北航接口] 总处理时间: {(datetime.now() - start_time).total_seconds():.2f}秒")
        
        return beihang_response

    except Exception as e:
        import traceback
        error_details = traceback.format_exc()
        logger.error(f"[北航接口] 处理过程中发生错误: {e}")
        logger.error(f"[北航接口] 错误详情: {error_details}")
        
        # 返回错误响应
        return {
            'routeInfos': [],
            'extraInfo': request.extraInfo if hasattr(request, 'extraInfo') else {}
        }

def validate_beihang_response_format(beihang_response: Dict[str, Any]) -> bool:
    """
    验证北航返回的数据格式是否正确
    """
    try:
        # 检查必需字段
        required_fields = ['routeInfos', 'extraInfo']
        for field in required_fields:
            if field not in beihang_response:
                logger.error(f"验证失败: 缺少必需字段 '{field}'")
                return False
        
        # 检查routeInfos格式
        route_infos = beihang_response['routeInfos']
        if not isinstance(route_infos, list):
            logger.error("验证失败: routeInfos 不是列表")
            return False
        
        for i, route_info in enumerate(route_infos):
            route_required_fields = ['outBizCode', 'routePoints', 'startTime', 'endTime', 'maxAlt', 'minAlt', 'width', 'length']
            for field in route_required_fields:
                if field not in route_info:
                    logger.error(f"验证失败: routeInfos[{i}] 缺少必需字段 '{field}'")
                    return False
            
            # 检查routePoints格式
            route_points = route_info['routePoints']
            if not isinstance(route_points, list):
                logger.error(f"验证失败: routeInfos[{i}].routePoints 不是列表")
                return False
            
            for j, point in enumerate(route_points):
                point_required_fields = ['longitude', 'latitude', 'altitude', 'speed', 'timestamp']
                for field in point_required_fields:
                    if field not in point:
                        logger.error(f"验证失败: routeInfos[{i}].routePoints[{j}] 缺少必需字段 '{field}'")
                        return False
        
        # 检查extraInfo格式
        extra_info = beihang_response['extraInfo']
        if not isinstance(extra_info, dict):
            logger.error("验证失败: extraInfo 不是字典")
            return False
        
        logger.info("北航数据格式验证通过!")
        return True
        
    except Exception as e:
        logger.error(f"验证过程中发生错误: {e}")
        return False

def visualize_environment_data(env_info: Dict[str, Any], config: Dict[str, Any]):
    """
    可视化环境数据（障碍物、禁飞区、敏感区域等）
    """
    try:
        from shapely.geometry import Point, Polygon, MultiPolygon, LineString
        from shapely import wkt
        from algorithm.coord_transform import PositionConvert
        
        fig, ax = plt.subplots(figsize=(14, 10))
        
        # 设置坐标转换器
        pos_converter = PositionConvert()
        ref_lat = config['ref_lat']
        ref_lon = config['ref_lon']
        
        # 收集所有坐标点用于自动调整范围
        all_x_coords = []
        all_y_coords = []
        
        # 绘制障碍物
        obstacles = env_info.get('obstacles', [])
        for i, obstacle_wkt in enumerate(obstacles):
            try:
                geom = wkt.loads(obstacle_wkt)
                if geom.geom_type == 'Polygon':
                    x, y = geom.exterior.coords.xy
                    # 转换为本地坐标
                    local_coords = []
                    for lon, lat in zip(x, y):
                        local_x, local_y = pos_converter.GPStoXY(lat, lon, ref_lat, ref_lon)
                        local_coords.append((local_x, local_y))
                        all_x_coords.append(local_x)
                        all_y_coords.append(local_y)
                    
                    if local_coords:
                        local_x_coords, local_y_coords = zip(*local_coords)
                        ax.fill(local_x_coords, local_y_coords, color='red', alpha=0.6, label=f'Obstacle {i+1}')
                elif geom.geom_type == 'MultiPolygon':
                    for j, poly in enumerate(geom.geoms):
                        x, y = poly.exterior.coords.xy
                        local_coords = []
                        for lon, lat in zip(x, y):
                            local_x, local_y = pos_converter.GPStoXY(lat, lon, ref_lat, ref_lon)
                            local_coords.append((local_x, local_y))
                            all_x_coords.append(local_x)
                            all_y_coords.append(local_y)
                        
                        if local_coords:
                            local_x_coords, local_y_coords = zip(*local_coords)
                            ax.fill(local_x_coords, local_y_coords, color='red', alpha=0.6, label=f'Obstacle {i+1}-{j+1}')
            except Exception as e:
                logger.error(f"绘制障碍物 {i+1} 时出错: {e}")
        
        # 绘制禁飞区
        no_fly_zones = env_info.get('noFlyZones', [])
        for i, zone_wkt in enumerate(no_fly_zones):
            try:
                geom = wkt.loads(zone_wkt)
                if geom.geom_type == 'Polygon':
                    x, y = geom.exterior.coords.xy
                    local_coords = []
                    for lon, lat in zip(x, y):
                        local_x, local_y = pos_converter.GPStoXY(lat, lon, ref_lat, ref_lon)
                        local_coords.append((local_x, local_y))
                        all_x_coords.append(local_x)
                        all_y_coords.append(local_y)
                    
                    if local_coords:
                        local_x_coords, local_y_coords = zip(*local_coords)
                        ax.fill(local_x_coords, local_y_coords, color='orange', alpha=0.4, label=f'No-Fly Zone {i+1}')
            except Exception as e:
                logger.error(f"绘制禁飞区 {i+1} 时出错: {e}")
        
        # 绘制敏感区域
        sensitive_areas = env_info.get('sensitiveAreas', [])
        for i, area in enumerate(sensitive_areas):
            try:
                if isinstance(area, dict) and 'geom' in area:
                    geom = wkt.loads(area['geom'])
                    area_type = area.get('type', 'Unknown')
                    
                    if geom.geom_type == 'Polygon':
                        x, y = geom.exterior.coords.xy
                        local_coords = []
                        for lon, lat in zip(x, y):
                            local_x, local_y = pos_converter.GPStoXY(lat, lon, ref_lat, ref_lon)
                            local_coords.append((local_x, local_y))
                            all_x_coords.append(local_x)
                            all_y_coords.append(local_y)
                        
                        if local_coords:
                            local_x_coords, local_y_coords = zip(*local_coords)
                            ax.fill(local_x_coords, local_y_coords, color='yellow', alpha=0.3, label=f'Sensitive Area {i+1} ({area_type})')
            except Exception as e:
                logger.error(f"绘制敏感区域 {i+1} 时出错: {e}")
        
        # 绘制任务区域边界（基于map_bounds）
        map_bounds = config.get('map_bounds', [-300, 0, 400, 400])
        task_bounds_x = [map_bounds[0], map_bounds[2], map_bounds[2], map_bounds[0], map_bounds[0]]
        task_bounds_y = [map_bounds[1], map_bounds[1], map_bounds[3], map_bounds[3], map_bounds[1]]
        ax.plot(task_bounds_x, task_bounds_y, 'g--', linewidth=2, label='Task Area Boundary')
        
        # 自动调整坐标轴范围
        if all_x_coords and all_y_coords:
            x_min, x_max = min(all_x_coords), max(all_x_coords)
            y_min, y_max = min(all_y_coords), max(all_y_coords)
            
            # 添加一些边距
            x_margin = (x_max - x_min) * 0.1
            y_margin = (y_max - y_min) * 0.1
            
            ax.set_xlim(x_min - x_margin, x_max + x_margin)
            ax.set_ylim(y_min - y_margin, y_max + y_margin)
            
            logger.info(f"环境数据坐标范围: X({x_min:.1f} ~ {x_max:.1f}), Y({y_min:.1f} ~ {y_max:.1f})")
            logger.info(f"任务区域边界: X({map_bounds[0]} ~ {map_bounds[2]}), Y({map_bounds[1]} ~ {map_bounds[3]})")
        
        # 设置图表属性
        ax.set_title("Beihang Environment Data Visualization", fontsize=14)
        ax.set_xlabel("X coordinate (m)")
        ax.set_ylabel("Y coordinate (m)")
        ax.grid(True, linestyle='--', alpha=0.6)
        ax.set_aspect('equal', adjustable='box')
        
        # 改进图例显示 - 放在图表外部
        handles, labels = ax.get_legend_handles_labels()
        if handles:
            by_label = dict(zip(labels, handles))
            ax.legend(by_label.values(), by_label.keys(), 
                     loc='center left', bbox_to_anchor=(1, 0.5),
                     fontsize=10, framealpha=0.8)
        
        # 调整布局防止图例被截断
        plt.tight_layout()
        plt.subplots_adjust(right=0.8)  # 为图例留出空间
        
        # 保存图片到visualization文件夹
        vis_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'visualization')
        os.makedirs(vis_dir, exist_ok=True)
        env_image_path = os.path.join(vis_dir, 'beihang_environment_visualization.png')
        plt.savefig(env_image_path, dpi=300, bbox_inches='tight')
        logger.info(f"北航环境数据可视化图片已保存到: {env_image_path}")
        
        plt.show()
        
    except Exception as e:
        logger.error(f"环境数据可视化失败: {e}")

def visualize_trajectories(trajectories: Dict[str, Any], time_paths: Dict[str, Any], config: Dict[str, Any]):
    """
    可视化规划结果轨迹
    """
    try:
        from client.visualization import plot_trajectories
        plot_trajectories({
            'trajectories': trajectories,
            'time_paths': time_paths
        }, config, save_image=False)
    except Exception as e:
        logger.error(f"轨迹可视化失败: {e}")

@app.get("/", include_in_schema=False)
def read_root():
    return {
        "message": "欢迎使用北航航线规划API",
        "endpoints": {
            "beihang_format": "/plan_route_beihang"
        },
        "description": "专门处理北航格式的无人机任务规划"
    }

# --- 启动服务 ---
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="北航航线规划API服务")
    parser.add_argument("--host", default="0.0.0.0", help="服务器主机地址")
    parser.add_argument("--port", type=int, default=8002, help="服务器端口")
    parser.add_argument("--visualize", action="store_true", help="启用可视化功能")
    
    args = parser.parse_args()
    
    # 设置全局可视化变量
    global_visualize_enabled = args.visualize
    
    print(f"启动北航航线规划API服务...")
    print(f"服务器地址: {args.host}:{args.port}")
    print(f"可视化功能: {'启用' if args.visualize else '禁用'}")
    
    # 检查端口是否被占用，自动处理冲突
    if is_port_in_use(args.port):
        print(f"警告: 端口 {args.port} 已被占用")
        print("正在自动尝试解决端口冲突...")
        
        if find_and_kill_process_on_port(args.port):
            print("成功杀死占用端口的进程")
            # 等待一下让端口释放
            import time
            time.sleep(2)
            
            # 再次检查端口是否已释放
            if is_port_in_use(args.port):
                print("警告: 端口仍然被占用，可能进程未能完全释放")
                print("请手动关闭占用该端口的程序后重试")
                sys.exit(1)
            else:
                print("端口已释放，继续启动服务")
        else:
            print("未能自动杀死占用端口的进程")
            print("请手动关闭占用该端口的程序后重试")
            sys.exit(1)
    
    print(f"API文档: http://{args.host}:{args.port}/docs")
    print(f"北航格式接口: http://{args.host}:{args.port}/plan_route_beihang")
    
    if args.visualize:
        print("\n注意: 可视化功能已启用，所有请求都会自动进行可视化")
    else:
        print("\n注意: 可视化功能已禁用，所有请求都不会进行可视化")
    
    uvicorn.run(app, host=args.host, port=args.port)